#!/usr/bin/env python
"""
GAmma-Ray FITting script based on MyGaRDiAn package.
Author: Zhaoqiang SHEN
"""
from __future__ import print_function
import os
import argparse
import multiprocessing
import numpy as np

# do not use multithreading, multiprocessing works better
os.environ['OMP_NUM_THREADS'] = '1'

def garLike(srcmap, mdlfile,
            healpix=False, slowTS=False,
            minbin=None, maxbin=None,
            centerSrc=None, ULTSLimit=10.,
            maskfile=None, maskext=None, maskType='flat',
            drmfile=None,
            outresult='results_garLike.dat', outmdl='model_fitted',
            outmcube=None, outcov=None,
            srcprf=None, prfmin=0.01, prfmax=100., prfbin=30, outprfpfx='chi2prf',
        ):
    l = NewLikeObj(None, None) # Manual Mode
    frees, fixs = l.parseModel(mdlfile)
    try:
        l.like = l._initChiSqObj(frees, fixs, srcmap, isHealpix=healpix, drmfile=drmfile)
    except TypeError:
        l.like = l._initChiSqObj(frees, fixs, srcmap, isHealpix=healpix)

    if not maskfile is None:
        l.like.addMask(mapFile=maskfile, extIdx=maskext, maskType=maskType)
    l._init_Minuit((minbin, maxbin))
    l.fit()

    if outcov:
        try:
            l.save_covariance(outcov)
        except Exception as e:
            print('[ERROR] fail to save covariance: %s'%(str(e)))
    if outresult:
        l.writeResults(outresult, centerSrc=centerSrc, ULTSLimit=ULTSLimit, slowTS=slowTS)
    if outmdl:
        l.saveModel(outmdl)
    if isinstance(outmcube, str):
        l.like.toModel(outmcube)

    if srcprf is not None:
        mypath_, myfile_ = os.path.split(outprfpfx)
        outprffile = os.path.join(mypath_,
                myfile_.split('.')[0]+'_%s.dat'%(''.join(srcprf.split())))
        l.profile(srcprf, norm_min=prfmin, norm_max=prfmax, outfile=outprffile, nbins=prfbin)

    print('.'*30, 'Done!', '.'*30)
    return l

def garSED(srcmap, mdlfile,
           healpix=False, slowTS=False,
           centerSrc=None, ULTSLimit=10.,
           maskfile=None, maskext=None, maskType='flat',
           drmfile=None,
           outresult='results_garLike.dat', outmdl='model_fitted', outmcube=None, outcov=None,
           njobs=None, ipart=None,
           srcprf=None, prfmin=0.01, prfmax=100., prfbin=50, outprfpfx='chi2prf',
        ):
    if (outresult is None) or ('%' in outresult):
        outresult_format = outresult
    else:
        mypath_, myfile_ = os.path.split(outresult)
        outresult_format = os.path.join(mypath_, myfile_.split('.')[0]+'_band%02i.dat')

    if (outmdl is None) or ('%' in outmdl):
        outmdl_format = outmdl
    else:
        mypath_, myfile_ = os.path.split(outmdl)
        outmdl_format = os.path.join(mypath_, myfile_.split('.')[0]+'_band%02i')

    if (outcov is None) or ('%' in outcov):
        outcov_format = outcov
    else:
        mypath_, myfile_ = os.path.split(outcov)
        outcov_format = os.path.join(mypath_, myfile_.split('.')[0]+'_band%02i.npz')

    l = NewLikeObj(None, None) # Manual Mode
    frees, fixs = l.parseModel(mdlfile)
    try:
        l.like = l._initChiSqObj(frees, fixs, srcmap, isHealpix=healpix, drmfile=drmfile)
    except TypeError:
        l.like = l._initChiSqObj(frees, fixs, srcmap, isHealpix=healpix)

    if not maskfile is None:
        l.like.addMask(mapFile=maskfile, extIdx=maskext, maskType=maskType)

    if (njobs is None) or (ipart is None):
        min_bin, max_bin = 0, l.like.data_nlayers
    else:
        assert njobs > ipart >= 0
        ebins = np.arange(l.like.data_nlayers-1, -1, -1)
        split_ebins = np.array_split(ebins, njobs)
        mypart = split_ebins[ipart]
        min_bin, max_bin = mypart.min(), mypart.max()+1
        print('[INFO] Calculate from bin [%s, %s] !' % (min_bin, max_bin-1))

    for ilayer in range(min_bin, max_bin):
        l._init_Minuit(ilayer)
        l.fit()
        if outcov_format:
            try:
                l.save_covariance(outcov_format%ilayer)
            except Exception as e:
                print('[ERROR] fail to save covariance for ilayer=%d: %s'%(ilayer,str(e)))

        if outresult_format:
            l.writeResults(outresult_format%ilayer, centerSrc=centerSrc,
                           ULTSLimit=ULTSLimit, slowTS=slowTS)
        if outmdl_format:
            l.saveModel(outmdl_format%ilayer)

        if srcprf is not None:
            mypath_, myfile_ = os.path.split(outprfpfx)
            outprffile = os.path.join(mypath_,
                    myfile_.split('.')[0]+'_band%02i_%s.dat'%(ilayer,''.join(srcprf.split())))
            l.profile(srcprf, norm_min=prfmin, norm_max=prfmax, outfile=outprffile, nbins=prfbin)

    if isinstance(outmcube, str):
        l.like.toModel(outmcube)

    print('.'*30, 'Done!', '.'*30)
    return l

def _garSED(*arg):
    garSED(*arg)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(prefix_chars='-+',
            description='The Command-Line Interface of MyGardian.MyLike.NewLikeObj')
    parser.add_argument("srcmap", type=str, help='srcmap')
    parser.add_argument("mdlfile", type=str,
            help='input model file, no bash variable allowed')
    parser.add_argument("-centerSrc", type=str, default=None,
            help='The name of center source')
    parser.add_argument("-ULTSLimit", type=float, default=10.,
            help='When TS < ULTSLimit, the UL of the center source will be calculated')
    parser.add_argument("-maskfile", type=str, default=None,
            help='The name of maskfile')
    parser.add_argument("-maskext", type=str, default=None,
            help='The extension index of maskfile')
    parser.add_argument("-masktype", type=str, default='flat', choices=['flat', 'structured'],
            help='The type of mask')
    parser.add_argument("-datfile", type=str, default='results_garLike.dat',
            help='The name of output dat file')
    parser.add_argument("-outmdl", type=str, default='model_fitted',
            help='The name of output model file (.xml|.yaml|.json)')
    parser.add_argument("-outmcube", type=str, default=None,
            help='The name of output mcube file (do not support multiprocessing)')
    parser.add_argument("-outcov", type=str, default=None,
            help='The name of output covariance data file)')
    parser.add_argument("+healpix", action='store_true', default=False,
            help='Whether the input map is healpix')
    parser.add_argument("+slowTS", action='store_true', default=False,
            help='Whether to get the accurate TS value')
    parser.add_argument("+sed", action='store_true', default=False,
            help='Whether to calculate sed')
    parser.add_argument("-njobs", type=int, default=1,
            help='The number of jobs you wish to spawn when calculating sed')
    parser.add_argument("-minbin", type=int, default=None,
            help='The bin range [minbin:maxbin] will be calculate in the global fit')
    parser.add_argument("-maxbin", type=int, default=None,
            help='The bin range [minbin:maxbin] will be calculate in the global fit')
    parser.add_argument("-drmfile", type=str, default=None,
            help='The name of the drm file')

    parser.add_argument("-srcprf", type=str, default=None,
            help='the name of src to evaluate the -2*lnL profile (def: None)')
    parser.add_argument("-prfmin", type=float, default=0.01,
            help='the min norm value in calc prof like (def: 0.01)')
    parser.add_argument("-prfmax", type=float, default=100.,
            help='the max norm value in calc prof like (def: 100.)')
    parser.add_argument("-prfbin", type=int, default=30,
            help='the numbers of norm values evaluated (def: 30)')
    parser.add_argument("-outprfpfx", type=str, default='chi2prf',
            help='the prefix of the output profile data file (def: chi2prf)')

    parser.add_argument("+glat2", action='store_true', default=False,
            help='whether to use glat2 (py2/py3), default: False (use MyGardian, py2 only)')

    args = parser.parse_args()

    if args.glat2:
        ## the new backend (support py2+py3)
        from glat2.likeold import NewLikeObj
    else:
        ## the old backend
        from MyGardian.MyLike import NewLikeObj

    if args.sed:
        if args.njobs > 1:
            pool = multiprocessing.Pool(processes=args.njobs)
            for ijob in range(args.njobs):
                err = pool.apply_async(_garSED, (
                        args.srcmap, args.mdlfile,
                        args.healpix, args.slowTS,
                        args.centerSrc, args.ULTSLimit,
                        args.maskfile, args.maskext, args.masktype,
                        args.drmfile,
                        args.datfile, args.outmdl, None, args.outcov,
                        args.njobs, ijob,
                        args.srcprf, args.prfmin, args.prfmax, args.prfbin,
                        args.outprfpfx,
                    )
                )
            pool.close()
            pool.join()
            err.get()
        else:
            garSED(args.srcmap, args.mdlfile,
                   args.healpix, args.slowTS,
                   centerSrc=args.centerSrc, ULTSLimit=args.ULTSLimit,
                   maskfile=args.maskfile, maskext=args.maskext,
                   maskType=args.masktype,
                   drmfile=args.drmfile,
                   outresult=args.datfile, outmdl=args.outmdl,
                   outmcube=args.outmcube, outcov=args.outcov,
                   srcprf=args.srcprf, prfmin=args.prfmin, prfmax=args.prfmax,
                   prfbin=args.prfbin, outprfpfx=args.outprfpfx)
    else:
        garLike(args.srcmap, args.mdlfile, healpix=args.healpix, slowTS=args.slowTS,
                minbin=args.minbin, maxbin=args.maxbin,
                centerSrc=args.centerSrc, ULTSLimit=args.ULTSLimit,
                maskfile=args.maskfile, maskext=args.maskext, maskType=args.masktype,
                drmfile=args.drmfile,
                outresult=args.datfile, outmdl=args.outmdl,
                outmcube=args.outmcube, outcov=args.outcov,
                srcprf=args.srcprf, prfmin=args.prfmin, prfmax=args.prfmax,
                prfbin=args.prfbin, outprfpfx=args.outprfpfx,
                )

